from PIL import Image
import numpy as np
import sys
import os

def rgb_to_yuv(r, g, b):
    # BT.601标准公式（标准范围）
    y = 16 + (65.481 * r + 128.553 * g + 24.966 * b) / 255
    u = 128 + (-37.797 * r - 74.203 * g + 112.0 * b) / 255
    v = 128 + (112.0 * r - 93.786 * g - 18.214 * b) / 255
    # 限制在标准范围内
    return np.clip(y, 16, 235), np.clip(u, 16, 240), np.clip(v, 16, 240)

def convert_to_yuv(image_path, output_path=None):
    img = Image.open(image_path).convert('RGBA')
    img_array = np.array(img)

    r, g, b, a = img_array[:, :, 0], img_array[:, :, 1], img_array[:, :, 2], img_array[:, :, 3]
    y, u, v = rgb_to_yuv(r, g, b)

    yuv_array = np.stack((y, u, v), axis=-1)
    yuv_image = Image.fromarray(yuv_array.astype('uint8'))

    if output_path is None:
        base_name = os.path.splitext(image_path)[0]
        output_path = f"{base_name}_yuv.png"
    yuv_image.save(output_path)
    print(f"YUV image saved to: {output_path}")

if __name__ == "__main__":
    if len(sys.argv) > 2:
        convert_to_yuv(sys.argv[1], sys.argv[2])
    elif len(sys.argv) == 2:
        convert_to_yuv(sys.argv[1])
    else:
        raise ValueError("Invalid number of arguments, Usage: python png_to_yuv.py <image_path>")
